/**
 \file CollinsKt.cxx
 Implementation of class CollinsKt.
 
 \author Thomas Burton 
 \date 2/28/12
 */

#include <algorithm>
#include <iostream>
#include <map>
#include <sstream>
#include <utility>

#include <boost/foreach.hpp>

#include <Math/ParamFunctor.h>
#include <TCanvas.h>
#include <TH1D.h>
#include <TMath.h>
#include <TRandom.h>

#include "tppmc/Bins.h"
#include "tppmc/CollinsKt.h"
#include "tppmc/DssFragmentation.h"
#include "tppmc/KeyedLists.h"
#include "tppmc/KkpParton.h"
#include "tppmc/parameters.h"

namespace tppmc {

namespace collins {

/**
 Function object evaluating Collins function.
 */
struct FunctorNew {
   
   // <pT^2> for the Gaussian pT distribution
   double mMeanPt2;
   // The unpolarised fragmentation function against which the positivity
   // bound is computed
   boost::shared_ptr<tppmc::FragmentationFunction> mUnpol;
   
   FunctorNew(boost::shared_ptr<tppmc::FragmentationFunction> unpol)
   : mMeanPt2(0.2)
   , mUnpol(unpol) {
   }
   
   enum EFunctorParameters {
      SCALE = 0,
      FAVOURED = 1,
      POLARISATION = 2
   };
   
   double Favoured(const Parton& parton, const Hadron& hadron,
                   double Q2) {
      return mUnpol->Evaluate(parton, hadron, Q2);
   }
   
   double Unfavoured(const Parton& parton, const Hadron& hadron,
                     double Q2) {
      return -Favoured(parton, hadron, Q2);
   }
   
   double operator()(double* point, double* parameters) {
      double z = point[0];
      double phi = point[1];
      double qs2 = pow(parameters[SCALE], 2.);
      int isFavoured = static_cast<int>(parameters[FAVOURED]);
      double depol_factor = parameters[POLARISATION];
      // The fragmentation functions are parametrized in terms of the
      // fragmentation function for a positive pion, and we choose the
      // favored or unfavored fragmentation function from this by
      // setting the quark type, up quark for favoured,
      // or down quark for unfavoured.
      double D1(0.);
      double collins(0.);
      // Favoured collins = D1(u-->pi+)
      // Unfavoured collins = -D1(d-->pi+) (note the minus sign)
      // For pi0, use average of favoured and unfavoured for that quark
      Hadron piplus;
      piplus.type = 211;
      piplus.z = z;
      if(0 == isFavoured) { // unfavoured
         D1 = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 1), piplus, qs2);
         collins = -D1;
      } // if
      else if(1 == isFavoured) {// favoured
         D1 = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 2), piplus, qs2);
         collins = D1;
      } // else if
      else if(2 == isFavoured) { // Mean of favoured and unfavoured
         // Based this on the DSS definition of pi0 fragmentation.
         // They use (q->pi0) = average of (q->pi+) and (q->pi-)
         // which is the same as (favoured->pi+) and (unfavoured->pi-)
         // with the charge conjugation assumptions.
         // e.g. their u->pi- is the same as d->pi+
         double favoured = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 2), piplus, qs2);
         double unfavoured = mUnpol->Evaluate(tppmc::Parton(0., 0., 0., 1), piplus, qs2);
         // Unpolarised is the average
         D1 = 0.5 * (favoured + unfavoured);
         // Unfavoured collins is negative of D1
         collins = 0.5 * (favoured - unfavoured);
      } // else if
     // Why is there no "spin" factor (sign) in this equation? Because what we want is the angular 
     // distribution referenced to the spin direction. This is accounted for properly 
     // when we calculate the angular distrubution of the hadrons in the fragmentation
     // because it is done with respect to the spin direction.
     // NOTE ALSO: The angular distribution that is given is with respect to the vector
     // p^ x jT^, so the dot product gives a cos(phi) behavior. Because this is perp. to
     // jt^ the hadron distribution will have a sin(phi) modulation.
      double full_FF = D1 + cos(phi) * depol_factor * collins;
      return std::max(full_FF, 0.); // Don't go negative!
   }
}; // FunctorNew

//
// class PositivityKt
//

void PositivityKt::Qa(std::string filename) {
   TCanvas canvas;
   canvas.Print((filename + "[").c_str());
   double depols[] = { 0.1 };
   UniformBins depolbins(1, 0.1, 0.9);
   mFunctions.Initialise(&depolbins, mQBins.get());
   std::ostringstream oss;
   std::cout << Partons().size() << " " << Hadrons().size() << std::endl;
   BOOST_FOREACH(int parton, Partons()) {
      std::cout << parton << std::endl;
      BOOST_FOREACH(int hadron, Hadrons()) {
         BOOST_FOREACH(double depol, depols) {
            Hadron thehadron;
            thehadron.type = hadron;
            thehadron.mParton.type = parton;
            thehadron.mParton.cos_spin_depol = depol;
            TF2* func = mFunctions.Find(thehadron, 2.);
            if(not func) {
               std::cout << "Buh?" << std::endl;
               continue;
            } // if
            oss << "Parton " << parton << " hadron " << hadron
            << " depol " << depol;
            std::cout << oss.str() << std::endl;
            TH1D hist("", oss.str().c_str(), 9, 0.05, 0.95);
            oss.str("");
            for(int i(1); i <= 9; ++i) {
               double integral = func->Integral(0.05, 0.05 + i * 0.1,
                                                -TMath::PiOver2(), TMath::PiOver2());
               hist.SetBinContent(i, integral);
            } // for
            hist.Scale(1. / hist.GetBinContent(hist.GetNbinsX()));
            hist.SetMinimum(0.);
            hist.Draw();
            canvas.Print(filename.c_str());
         } // BOOST_FOREACH
      } // BOOST_FOREACH
   } // BOOST_FOREACH
   canvas.Print((filename + "]").c_str());
}

PositivityKt::PositivityKt()
: mMeanPt2(0.2) {
}

PositivityKt::PositivityKt(FragFcnPtr unpol, double meanPt2,
                           const std::set<int>& hadrons)
: mMeanPt2(meanPt2)
, mUnpol(unpol) {
   mHadrons.assign(hadrons.begin(), hadrons.end());
}

bool PositivityKt::Initialise() {
   if(not mUnpol.get()) {
      mUnpol = FragFcnPtr(new DssFragmentation(LO));
   } // if
   mUnpol->Initialise();
   mQBins.reset(new UniformBins(230, sqrt(MinimumScale()),
                                sqrt(MaximumScale())));
   UniformBins depol(20, 0., 1.);
   mFunctions.Initialise(&depol, mQBins.get());
   // Populate default hadron list if not set via constructor.
   if(mHadrons.empty()) {
      int hadrons[3] = { -211, 111, 211 };
      mHadrons.assign(hadrons, hadrons + 3);
   } // if
   int partons[2] = {2, 1};
   mPartons.assign(partons, partons + 2);
   return mQBins.get();
}

const std::vector<int>& PositivityKt::Hadrons() const {
   return mHadrons;
}

const std::vector<int>& PositivityKt::Partons() const {
   return mPartons;
}

PositivityKt::~PositivityKt() { // Nada - only smart pointers are used
}

void PositivityKt::ComputeFragmentationProbabilities() {
   std::vector<double> QBins = mQBins->Edges();
   std::vector<int> hadronlist = Hadrons();
   std::vector<int> partonlist = Partons();
   BOOST_FOREACH(int parton, partonlist) {
      FragTablePtr table(new FragmentationTable(mQBins->Edges(),
                                                hadronlist));
      mProbabilities.insert(std::make_pair(parton, table));
      // Loop over hadrons and Q bins and integrate the function
      BOOST_FOREACH(int hadron, hadronlist) {
         Hadron part;
         part.type = hadron;
         part.mParton.type = parton;
         // Overall probabilities of hadron production don't depend on
         // polarisation, so use zero
         part.mParton.cos_spin_depol = 0.;
         for(int i(0); i < mQBins->Number(); ++i) {
            double Q = QBins.at(i); // Lower edge of bin
            TF2* function = mFunctions.Find(part, Q);
            assert(function);
            double integral = function->Integral(ZMin(), ZMax(),
                                                 -TMath::Pi(), TMath::Pi());
            mProbabilities[parton]->Set(hadron, Q, integral);
         } // for
      } // BOOST_FOREACH
      mProbabilities[parton]->Normalise(1.);
   } // BOOST_FOREACH
   
}

int PositivityKt::ThrowHadronType(const Parton& parton, double Q2) {
   if(INVALID_HADRON not_eq mHadron) { // Force to pick certain type
      return mHadron;
   } // if
   if(mProbabilities.empty()) { // Lazy initialisation
      ComputeFragmentationProbabilities();
   } // if
   FragTablePtr table = mProbabilities[parton.type];
   /**\todo Need protection against invalid result */
   // Passed Q2 but table stores probabilities in bins in Q --> sqrt
   // Don't allow it to go below the minimum scale
   double Q = std::max(sqrt(Q2), sqrt(mUnpol->MinimumScale()));
   std::pair<int, bool> hadron = table->Random(Q);
   if(not hadron.second) {
      std::cerr << "Threw invalid hadron type... using pi0" << std::endl;
      hadron.first = 111; // PIZERO;
   } // if
   return hadron.first;
}

std::vector<Hadron> PositivityKt::Fragment(const Parton& parton,
                                           double scale) {
   if(not Supports(parton)) {
      return mUnpol->Fragment(parton, scale);
   } // if
   Hadron hadron;
   hadron.type = ThrowHadronType(parton, scale);
   switch(hadron.type) {
      case  211:
      case -211: // Deliberate fall-through
      case  111: // Deliberate fall-through
         break;
      default:
         std::cerr << "Threw hadron type " << hadron.type << " for parton " << 
         parton << " at scale " << scale << "? Returning DSS instead" << 
         std::endl;
         return mUnpol->Fragment(parton, scale);
   } // switch(h)
   hadron.mParton = parton;
   // Find the function with the parameters for this hadron/parton/scale
   // Input scale is Q2 but functions are stored in bins in Q
   TF2* function = mFunctions.Find(hadron, sqrt(scale));
   hadron.favoured = function->GetParameter(FunctorNew::FAVOURED);
   if(not function) {
      std::cerr << "No function for hadron " << hadron.type <<
      " parton " << hadron.mParton.type << " depol " <<
      hadron.mParton.cos_spin_depol << " at scale " << sqrt(scale)
      <<std::endl;
      assert(function);
   } // if
     // Randomly throw z and phi according to the function
     // Don't use TF2::SetRange() to restrict the z range to throw, as this
     // clears the integral cache. Instead keep throwing until until we
     // get a value in the required range.
   hadron.z = -1.;
   while(hadron.z > mZUpper or hadron.z < mZLower) {
      function->GetRandom2(hadron.z, hadron.phigen);
   } // while
     // Randomly generate kt.
   double jt_sigma = sqrt(mMeanPt2/2.0);  
   double jt; 
   // Throw jt according to an exponential distribution such
   // that mean jt^2 is 0.2
   do {
      jt = jt_sigma*sqrt( -2.0*log(gRandom->Uniform(1.0)) ); 	  
   } while(jt>(3.0*jt_sigma));
   TVector3 mom(jt, 0., 0.);
   // TVector3::SetPhi() rotates phi,
   // leaving magnitude and theta unchanged.
   mom.SetPhi(hadron.phigen);
   hadron.px = mom.Px();
   hadron.py = mom.Py();
   hadron.pz = parton.PxPyPz().Mag() * hadron.z;
   hadron.mFrag = Evaluate(parton, hadron, scale);
   // Find the Collins component by evaluating the fragmentation function
   // with zero polarisation and subtracting from the value obtained
   // above.
   Hadron unpolarised(hadron);
   unpolarised.mParton.cos_spin_depol = 0.;
   hadron.mCollins = hadron.mFrag - Evaluate(parton, unpolarised, scale);
   return std::vector<Hadron>(1, hadron);
}

double PositivityKt::Evaluate(const tppmc::pythpart& parton,
                              const Hadron& hadron,
                              double scale) {
   if(not Supports(parton)) {
      return mUnpol->Evaluate(parton, hadron, scale);
   } // if
   // Technically if the hadron probability computations went OK, we
   // shouldn't be able to throw anything other than pions, but in case
   // something weird happened return the unpolarised result.
   switch(hadron.Type()) {
      case  211:
      case -211: // Deliberate fall-through
      case  111: // Deliberate fall-through
         break;
      default:
         return mUnpol->Evaluate(parton, hadron, scale);
         break;
   } // switch(h)
   // Get the function with parameters appropriate for this hadron/parton/
   // scale (remember the functions are in bins in Q, not Q2)
   TF2* function = mFunctions.Find(hadron, sqrt(scale));
   assert(function);
   return function->Eval(hadron.z,
                         hadron.PxPyPz().Phi(),
                         hadron.PxPyPz().Pt());
}

bool PositivityKt::Supports(const Parton& parton) const {
   return std::find(mPartons.begin(), mPartons.end(),
                    parton.Pdg()) not_eq mPartons.end();
}

std::string PositivityKt::Name() const {
   std::ostringstream oss;
   oss << "DSS + Collins, kT-dependent positivity bound, <pt2> = " <<
   mMeanPt2;
   return oss.str();
}

/**
 Returns true if the fragmentation of the input parton type
 is supported by this function.
 */
bool PositivityKt::Supports(const Hadron& hadron) const {
   return hadron.type not_eq 111 and abs(hadron.type) not_eq 211;
}

void PositivityKt::SetHadron(int i) {
   FragmentationFunction::SetHadron(i);
   mUnpol->SetHadron(i);
}

void PositivityKt::SetZRange(double lower, double upper) {
   FragmentationFunction::SetZRange(lower, upper);
   mUnpol->SetZRange(lower, upper);
}

double PositivityKt::ZMin() const {
   return mUnpol->ZMin();
}

double PositivityKt::ZMax() const {
   return mUnpol->ZMax();
}

double PositivityKt::MinimumScale() const {
   return mUnpol->MinimumScale();
}

double PositivityKt::MaximumScale() const {
   return mUnpol->MaximumScale();
}

//
// class FunctionTable
//

FunctionTable::FunctionTable()
: mQ(NULL) {
}

FunctionTable::~FunctionTable() {
   // Delete the functors we created for each ROOT function.
   // Note f is a reference to a pointer so we can modify it to NULL
   // after deletion
   BOOST_FOREACH(FunctorNew*& f, mFunctors) {
      if(f) {
         delete f;
         f = NULL;
      } // if
   } // BOOST_FOREACH
}

void FunctionTable::Speak() {
   typedef Table::value_type Entry;
   BOOST_FOREACH(Entry& entry, mFunctions) {
      entry.first.Speak();
      for(int i(1); i <= mQ->GetNbins(); ++i) {
         std::cout << "\t(bin Q = " << mQ->GetBinLowEdge(i) << ") " <<
         entry.second.at(i - 1)->GetName() << " " << 
         entry.second.at(i - 1)->GetTitle() << std::endl;
      } // for
   } // BOOST_FOREACH
}

TF2* FunctionTable::Find(const Hadron& hadron, double Q) {
   // Determine which Q bin was requested.
   // If Q is outside the range covered, use the first or last bin
   // as appropriate
   int bin = mQ->FindBin(Q);
   if(bin < 1) { // Too low!
      bin = 1;
   } // if
   else if(bin > mQ->GetNbins()) { // Too high!
      bin = mQ->GetNbins();
   }
   bin -= 1; // Different ranges for TAxis [1,N] and vector [0,N)
   typedef Table::value_type Entry;
   BOOST_FOREACH(Entry& entry, mFunctions) {
      // Entry.first is the FragBin object
      // Entry.second is the vector<TF2*> with Q bins
      if(entry.first.Match(hadron)) { // Found the hadron/parton/depol
         return entry.second.at(bin); // Get the TF2* for this Q bin
      } // if
   } // BOOST_FOREACH
   return NULL;
}

TF2* FunctionTable::BuildFunction(FragBin bin, double Q) {
   // I would prefer to pass the functor by value, but I *think* that is
   // the source of some random crashing (ROOT does love it's pointers).
   // At least, I haven't found another source for the crash, so until I
   // do this should stay as a pointer.
   // The FunctorNews are collected in a list and deleted in the
   // destructor.
   mFunctors.push_back(new FunctorNew(mUnpol));
   // Despite what the ROOT documentation says, it seems TF2 needs the
   // class name passed as an argument when using a functor.
   // We also need to make sure the function has a unique "ROOT name"
   // otherwise it seems to crash like a m****r f****r, at least with
   // some ROOT versions.
   std::ostringstream oss;
   oss << "positivity_kt_fcn_" << bin.favoured << "_" << bin.depolMin
   << "_" << bin.depolMax << "_" << mQ->FindBin(Q) - 1;
   TF2* f = new TF2(oss.str().c_str(),
                    mFunctors.back(),
                    0.05, 1., -TMath::Pi(), TMath::Pi(), 3,
                    "tppmc::collins::FunctorNew");
   f->SetParameter(FunctorNew::SCALE, Q);
   f->SetParameter(FunctorNew::FAVOURED, bin.favoured);
   f->SetParameter(FunctorNew::POLARISATION, bin.depolMin);
   // Provide a useful title describing the parameters
   oss.str("");
   oss << mUnpol->Name() << " for Q = " << Q
   << " favoured = " << bin.favoured << " polarisation = "
   << bin.depolMin;
   f->SetTitle(oss.str().c_str());
   return f;
}

/** \todo What about multiple initialisations? */
int FunctionTable::Initialise(const Bins* depolbins, const Bins* Qbins) {
   mUnpol = FragFcnPtr(new tppmc::DssFragmentation);
   typedef std::vector<std::pair<double, double> > MinMaxVec;
   // Get the (min, max) of each bin, in both Q and depolarisation
   MinMaxVec depolBinsMinMax = depolbins->MinMax();
   MinMaxVec QBinMinMax = Qbins->MinMax();
   // Initialise the TAxis with the bin edges
   // We use this to convert Q values into bin number
   mQ = new TAxis(Qbins->Number(), &(Qbins->Edges().at(0)));
   typedef MinMaxVec::value_type MinMax;
   // Vector of NULL TF2 pointers to initialise each entry
   TF2Vec prototype(Qbins->Number(), (TF2*)(NULL));
   BOOST_FOREACH(MinMax d, depolBinsMinMax) {
      BOOST_FOREACH(MinMax Q, QBinMinMax) {
         for(int k(0); k < 3; ++k) { // unfavoured, favoured
            FragBin bin;
            bin.depolMin = d.first;
            bin.depolMax = d.second;
            bin.favoured = k;
            if(mFunctions.find(bin) == mFunctions.end()) {
               mFunctions.insert(std::make_pair(bin, prototype));
            } // if
              // Subtract 1 from bin in TAxis due to range [1, N]
            TF2* f = BuildFunction(bin, Q.first);
            mFunctions[bin].at(mQ->FindBin(Q.first) - 1) = f;
         } // for
      } // for
   } // for
   return 0;
}

//
// class KtDependent
//

std::string KtDependent::Find(const Hadron& hadron) const {
   typedef Table::value_type Pair;
   BOOST_FOREACH(Pair p, mTable) {
      if(p.first.Match(hadron)) return p.second;
   }
   return "";
}

} // namespace collins

} // namespace tppmc
